import numpy as np
from scipy.spatial.transform import Rotation
from scipy.linalg import logm


def trans_error(trans1, trans2):
    """
    计算两个平移向量之间的误差
    
    参数:
    trans1 (3x1 numpy array): 第一个平移向量
    trans2 (3x1 numpy array): 第二个平移向量
    
    返回:
    trans_error (float): 平移向量之间的欧几里得距离
    """
    return np.linalg.norm(trans1 - trans2)


def rot_error(rot1, rot2):
    """
    计算两个旋转矩阵之间的误差
    
    参数:
    rot1 (3x3 numpy array): 第一个旋转矩阵
    rot2 (3x3 numpy array): 第二个旋转矩阵
    
    返回:
    rot_error (float): 旋转矩阵之间的对数映射误差
    """
    error_matrix = rot1.T @ rot2
    
    angle_vec = Rotation.from_matrix(error_matrix).as_rotvec()
    return np.linalg.norm(angle_vec)


def pose_error(pose1, pose2):
    """
    计算两个位姿之间的误差
    
    参数:
    pose1 (list): [trans1, rot1] 第一个位姿,其中
        trans1 (3x1 numpy array): 平移向量
        rot1 (3x3 numpy array): 旋转矩阵
    pose2 (list): [trans2, rot2] 第二个位姿,其中
        trans2 (3x1 numpy array): 平移向量
        rot2 (3x3 numpy array): 旋转矩阵
        
    返回:
    trans_error (float): 平移向量之间的欧几里得距离
    rot_error (float): 旋转矩阵之间的对数映射误差
    """
    # 提取位姿
    trans1, rot1 = pose1
    trans2, rot2 = pose2

    return trans_error(trans1, trans2), rot_error(rot1, rot2)


if __name__ == "__main__":
    pose1 = [np.array([1, 2, 3]), np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])]
    pose2 = [np.array([4, 5, 6]), np.array([[0.5, 0.5, 0.5], [-0.5, 0.5, -0.5], [-0.5, -0.5, 0.5]])]

    t_error, R_error = pose_error(pose1, pose2)
    print(f"Translation error: {t_error:.4f}")
    print(f"Rotation error: {R_error:.4f}")